package com.smartcampusbackend.service.impl;

import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.net.URI;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.text.SimpleDateFormat;
import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.Base64;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.JsonNode;

public class SparkWebSocketClient {

    public static String sendAndReceive(String hostUrl, String appid, String apiKey, String apiSecret, String question) {
        try {
            // 1. 构造带签名的ws地址
            String authUrl = getAuthUrl(hostUrl, apiKey, apiSecret);

            // 2. 构造WebSocket客户端
            CountDownLatch latch = new CountDownLatch(1);
            StringBuilder answer = new StringBuilder();

            WebSocketClient client = new WebSocketClient(new URI(authUrl)) {
                @Override
                public void onOpen(ServerHandshake handshakedata) {
                    // 3. 发送请求数据
                    try {
                        Map<String, Object> payload = buildRequestPayload(appid, question);
                        ObjectMapper mapper = new ObjectMapper();
                        String json = mapper.writeValueAsString(payload);
                        this.send(json);
                    } catch (Exception e) {
                        e.printStackTrace();
                        latch.countDown();
                    }
                }

                @Override
                public void onMessage(String message) {
                    System.out.println("AI返回原始消息: " + message); // 新增日志
                    try {
                        ObjectMapper mapper = new ObjectMapper();
                        JsonNode root = mapper.readTree(message);
                        int code = root.path("header").path("code").asInt();
                        if (code != 0) {
                            answer.append("接口返回错误: ").append(root.path("header").path("message").asText());
                            latch.countDown();
                            return;
                        }
                        JsonNode choices = root.path("payload").path("choices");
                        JsonNode textArray = choices.path("text");
                        if (textArray.isArray()) {
                            for (JsonNode textItem : textArray) {
                                System.out.println("textItem: " + textItem.toString());
                                answer.append(textItem.path("content").asText());
                            }
                        }
                        int status = root.path("header").path("status").asInt();
                        if (status == 2) {
                            latch.countDown();
                        }
                    } catch (Exception e) {
                        answer.append("解析AI返回内容出错");
                        latch.countDown();
                    }
                }

                @Override
                public void onClose(int code, String reason, boolean remote) {
                    latch.countDown();
                }

                @Override
                public void onError(Exception ex) {
                    answer.append("WebSocket连接出错: ").append(ex.getMessage());
                    latch.countDown();
                }
            };

            client.connect();
            // 等待AI返回，最长30秒
            latch.await(30, TimeUnit.SECONDS);
            client.close();
            return answer.toString();

        } catch (Exception e) {
            e.printStackTrace();
            return "AI服务暂时不可用，请稍后再试。";
        }
    }

    // 构造请求体
    private static Map<String, Object> buildRequestPayload(String appid, String question) {
        Map<String, Object> payload = new HashMap<>();
        Map<String, Object> header = new HashMap<>();
        header.put("app_id", appid);

        Map<String, Object> parameter = new HashMap<>();
        Map<String, Object> chat = new HashMap<>();
        chat.put("domain", "lite");
        chat.put("temperature", 0.5);
        chat.put("max_tokens", 2048);
        chat.put("top_k", 4);
        parameter.put("chat", chat);

        Map<String, Object> payloadData = new HashMap<>();
        List<Map<String, Object>> textList = new ArrayList<>();
        Map<String, Object> textItem = new HashMap<>();
        textItem.put("role", "user");
        textItem.put("content", question);
        textList.add(textItem);

        Map<String, Object> message = new HashMap<>();
        message.put("text", textList);
        payloadData.put("message", message);

        payload.put("header", header);
        payload.put("parameter", parameter);
        payload.put("payload", payloadData);

        return payload;
    }

    // 构造带签名的ws地址
    private static String getAuthUrl(String hostUrl, String apiKey, String apiSecret) throws Exception {
        URI uri = new URI(hostUrl);
        String host = uri.getHost();
        String date = getGMTDate();
        String algorithm = "hmac-sha256";
        String headers = "host date request-line";
        String requestLine = "GET " + uri.getPath() + " HTTP/1.1";

        String signatureOrigin = "host: " + host + "\n" +
                "date: " + date + "\n" +
                requestLine;
        Mac mac = Mac.getInstance("HmacSHA256");
        mac.init(new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "HmacSHA256"));
        byte[] signData = mac.doFinal(signatureOrigin.getBytes(StandardCharsets.UTF_8));
        String signature = Base64.getEncoder().encodeToString(signData);

        String authorizationOrigin = String.format(
                "api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"",
                apiKey, algorithm, headers, signature
        );
        String authorization = Base64.getEncoder().encodeToString(authorizationOrigin.getBytes(StandardCharsets.UTF_8));

        String url = hostUrl + "?authorization=" + URLEncoder.encode(authorization, "UTF-8")
                + "&date=" + URLEncoder.encode(date, "UTF-8")
                + "&host=" + URLEncoder.encode(host, "UTF-8");
        return url;
    }

    // 获取GMT时间
    private static String getGMTDate() {
        SimpleDateFormat sdf = new SimpleDateFormat("E, dd MMM yyyy HH:mm:ss z", Locale.US);
        sdf.setTimeZone(TimeZone.getTimeZone("GMT"));
        return sdf.format(new Date());
    }


}
